# 1. 安装必要的库（在命令行运行一次）：
# pip install datasets scikit-learn pandas matplotlib openpyxl

import pandas as pd
from datasets import load_dataset
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
import matplotlib.pyplot as plt
import numpy as np

# 2. 加载 Stanford Sentiment Treebank 二分类版（binary SST）
#    参数 'default' 对应二分类（0=neg,1=pos）

sst2 = load_dataset('glue', 'sst2', trust_remote_code=True)

train = sst2['train']

# 手动二值化：大于 0.5 视为正面，<= 0.5 视为负面
# train_labels = [1 if score > 0.5 else 0 for score in train_data['label']]
# train_texts = train_data['sentence']
train_texts  = train['sentence']
train_labels = train['label']  # 0 或 1

# 文本向量化
vectorizer = CountVectorizer()
X_train = vectorizer.fit_transform(train_texts)

# 4. 训练 Multinomial Naive Bayes
clf = MultinomialNB()
clf.fit(X_train, train_labels)

# 5. 读取短视频评论（Excel 文件）
#    假设评论都在第 1 列，没有表头
df_comments = pd.read_excel('情感分析数据.xlsx', header=None)
comments = df_comments.iloc[:, 0].astype(str).tolist()

# 6. 对评论文本做同样的向量化并分类预测
X_comments = vectorizer.transform(comments)
preds = clf.predict(X_comments)

# 7. 统计积极／消极评论数与比例
positive_count = np.sum(preds == 1)
negative_count = np.sum(preds == 0)
total = positive_count + negative_count

print(f"积极情感评论数: {positive_count}，占比: {positive_count/total:.2%}")
print(f"消极情感评论数: {negative_count}，占比: {negative_count/total:.2%}")

# 8. 绘制饼状图
labels = ['积极', '消极']
counts = [positive_count, negative_count]

# 指定中文字体，SimHei 是黑体，Windows 上一般都有；Linux/macOS 也可能已安装
plt.rcParams['font.sans-serif'] = ['SimHei']
# 解决负号 '-' 显示成方块的问题
plt.rcParams['axes.unicode_minus'] = False
plt.figure(figsize=(6,6))
plt.pie(counts, labels=labels, autopct='%1.1f%%')
plt.title('短视频评论情感倾向分布')
plt.axis('equal')  # 保持饼图为圆形
plt.show()
